#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @File     : Modules.py
# @Time     : 2024/8/14 14:54


import torch
import torch.nn as nn
import torch.nn.functional as F


class ScaledDotProductAttention(nn.Module):
    """缩放点积注意力"""

    def __init__(self, temperature, attn_dropout=0.1):
        self.temperature = temperature
        self.dropout = nn.Dropout(attn_dropout)

    def forward(self, q, k, v, mask=None):
        # Q * K^T / sqrt(d_K)
        attn = torch.matmul(q / self.temperature, k.transpose(2, 3))

        if mask is not None:
            print("Mask ScaledDotProductAttention")
            attn = attn.masked_fill(mask == 0, -1e-9)

        # softmax(Q * K^T / sqrt(d_K))
        attn = self.dropout(F.softmax(attn, dim=-1))

        # softmax(Q * K^T / sqrt(d_K)) * V
        output = torch.matmul(attn, v)

        return output, attn